﻿Imports System.Runtime.InteropServices

Public Class DeviceSleepManager
  Implements IDisposable

  ' Constants
  Private Const PPN_UNATTENDEDMODE As Integer = 3
  Private Const POWER_STATE_UNATTEND As UInteger = &H400000
  Private Const POWER_STATE_BACKLIGHTON As UInteger = &H2000000
  Private Const POWER_NAME As UInteger = 1
  Private Const POWER_FORCE As UInteger = &H1000

  Private Enum DevicePowerState As Integer
    Unspecified = -1
    D0 = 0 'Full On: full power, full functionality
    D1     'Low Power On: fully functional at low power/performance
    D2     'Standby: partially powered with automatic wake
    D3     'Sleep: partially powered with device initiated wake
    D4     ' Off: unpowered
  End Enum

  ' Interop functions
  Private Declare Function PowerPolicyNotify Lib "coredll.dll" (ByVal dwMessage As Integer, ByVal dwData As Integer) As Boolean
  Private Declare Function RequestPowerNotifications Lib "coredll.dll" (ByVal hMsgQueue As IntPtr, ByVal flags As UInteger) As IntPtr
  Private Declare Function StopPowerNotifications Lib "coredll.dll" (ByVal hHandle As IntPtr) As Integer
  Private Declare Function CreateMsgQueue Lib "coredll.dll" (ByVal name As String, ByRef options As MsgQOptions) As IntPtr
  Private Declare Function CloseMsgQueue Lib "coredll.dll" (ByVal msgqueue As IntPtr) As Boolean
  Private Declare Function WaitForSingleObject Lib "coredll.dll" (ByVal hHandle As IntPtr, ByVal wait As Integer) As UInteger
  Private Declare Function ReadMsgQueue Lib "coredll.dll" (ByVal msgqueue As IntPtr, ByVal lpBuffer() As Byte, ByVal bufsize As UInteger, ByRef numRead As UInteger, ByVal timeout As Integer, ByRef flags As UInteger) As Boolean
  Private Declare Sub SystemIdleTimerReset Lib "coredll.dll" ()
  Private Declare Function SetPowerRequirement Lib "coredll.dll" (ByVal pvDevice As String, ByVal DeviceState As Integer, ByVal DeviceFlags As UInteger, ByVal pvSystemState As IntPtr, ByVal Reserved As ULong) As IntPtr
  Private Declare Function ReleasePowerRequirement Lib "coredll.dll" (ByVal hPowerReq As IntPtr) As UInteger

  ' MessageQueue structure options
  <StructLayout(LayoutKind.Sequential)> _
  Public Structure MsgQOptions
    Public dwSize As UInteger
    Public dwFlags As UInteger
    Public dwMaxMessages As UInteger
    Public cbMaxMessage As UInteger
    Public bReadAccess As Boolean
  End Structure

  ' Private variables
  Private _guid As Guid = Guid.NewGuid
  Private _powerQueue As String = String.Empty
  Private _msgQueueHandle As IntPtr = IntPtr.Zero
  Private _powerNotificationHandle As IntPtr = IntPtr.Zero
  Private _messageLoopThread As Threading.Thread = Nothing
  Private _loopUntil As Boolean = False
  Private _unattendEnabled As Boolean = False
  Private _isPaused As Boolean = False
  Private _hWavPowerReq As IntPtr = IntPtr.Zero

  Public Sub New()
  End Sub

  Public Sub StartPreventDeviceSleep(Optional ByVal msgQueueSalt As String = "")
    If Not Me._unattendEnabled Then
      Me._unattendEnabled = True
      Me.WriteErrorOnMainThread("DeviceSleepManager :: StartPreventDeviceSleep")

      ' Set unattended mode to enabled
      PowerPolicyNotify(PPN_UNATTENDEDMODE, True)

      ' Set WAV device to play while device is in unattended mode
      Me.SetWavDevicePowerRequirement()

      ' Create messagequeue name
      Me._powerQueue = msgQueueSalt & Me._guid.ToString("N")

      ' Set the options up for the messagequeue
      Dim mqopt As New MsgQOptions
      mqopt.dwFlags = 0
      mqopt.dwMaxMessages = 20
      mqopt.cbMaxMessage = 50
      mqopt.bReadAccess = True
      mqopt.dwSize = System.Runtime.InteropServices.Marshal.SizeOf(mqopt)

      ' Create the messagequeue
      Me._msgQueueHandle = CreateMsgQueue(Me._powerQueue, mqopt)

      ' Create the power notification handle
      Me._powerNotificationHandle = RequestPowerNotifications(Me._msgQueueHandle, 1)

      ' Create a thread that will monitor our message queue
      Me._messageLoopThread = New Threading.Thread(AddressOf ThreadMsgQueueLoop)
      Me._messageLoopThread.Name = "DeviceSleepMessageLoopThread"
      SyncLock MainThreadPool
        MainThreadPool.Add(Me._messageLoopThread)
      End SyncLock
      Me._messageLoopThread.Start()
    End If
  End Sub

  Public Sub StopPreventDeviceSleep()
    If Me._unattendEnabled Then
      Me._unattendEnabled = False
      Me.WriteErrorOnMainThread("DeviceSleepManager :: StopPreventDeviceSleep")

      ' Release the requirement that the WAV device remain on while device is in unattended mode
      Me.ReleaseWavDevicePowerRequirement()

      ' Set unattended mode to disabled
      PowerPolicyNotify(PPN_UNATTENDEDMODE, False)

      If Me._messageLoopThread IsNot Nothing Then
        ' Set the loopuntil variable to false so that the while loop no longer executes
        Me._loopUntil = False

        ' Abort the message loop monitoring thread
        Me._messageLoopThread.Abort()
      End If

      ' Inform WM that we no longer need power notifications
      If Me._powerNotificationHandle <> IntPtr.Zero Then
        StopPowerNotifications(Me._powerNotificationHandle)
      End If

      ' Destroy the message queue
      If Me._msgQueueHandle <> IntPtr.Zero Then
        CloseMsgQueue(Me._msgQueueHandle)
      End If
    End If
  End Sub

  Public Property PausePreventDeviceSleep() As Boolean
    Get
      Return Me._isPaused
    End Get
    Set(ByVal value As Boolean)
      Me._isPaused = value

      If Me._isPaused Then
        Me.WriteErrorOnMainThread("DeviceSleepManager :: PausePreventDeviceSleep :: Start Pause")
        Me.ReleaseWavDevicePowerRequirement()
      Else
        Me.WriteErrorOnMainThread("DeviceSleepManager :: StopPreventDeviceSleep :: End Pause")
        Me.SetWavDevicePowerRequirement()
      End If
    End Set
  End Property

  Private Sub ThreadMsgQueueLoop()
    Try
      Dim buf(49) As Byte ' 50 element array of bytes
      Dim numberRead As UInteger = 0
      Dim flags As UInteger = 0
      Dim isUnattended As Boolean = False

      ' Loop until this is false or the thread is aborted
      Me._loopUntil = True
      While Me._loopUntil
        ' Block until a notification is put on the messagequeue
        If WaitForSingleObject(Me._msgQueueHandle, 4000) = 0 Then
          ' Read the latest item(s) off the messagequeue
          ReadMsgQueue(Me._msgQueueHandle, buf, buf.Length, numberRead, -1, flags)

          ' Convert the binary data in buf into an unsigned integer
          Dim curState As UInteger = ConvertByteArray(buf, 4)

          If (curState And POWER_STATE_UNATTEND) = POWER_STATE_UNATTEND Then
            ' Screen turned off!
            isUnattended = True
            Me.WriteErrorOnMainThread("DeviceSleepManager :: ThreadMsgQueueLoop :: Screen Off")
          ElseIf (curState And POWER_STATE_BACKLIGHTON) = POWER_STATE_BACKLIGHTON Then
            ' Screen turned on!
            isUnattended = False
            Me.WriteErrorOnMainThread("DeviceSleepManager :: ThreadMsgQueueLoop :: Screen On")
          End If
        End If

        If isUnattended AndAlso Not Me._isPaused Then
          SystemIdleTimerReset()
        End If
      End While
    Catch ex As Threading.ThreadAbortException
      ' Ignore thread abort exceptions since they are thrown when the thread is intentionally being terminated
      Me.StopPreventDeviceSleep()
    Catch ex As Exception
      ' Nothing to do for other exceptions (yet)
      Me.StopPreventDeviceSleep()
    Finally
      SyncLock MainThreadPool
        MainThreadPool.Remove(Me._messageLoopThread)
      End SyncLock
      Me._messageLoopThread = Nothing
    End Try
  End Sub

  ' Utility function to convert 4 adjacent bytes into an unsigned integer
  Private Shared Function ConvertByteArray(ByVal array() As Byte, ByVal offset As Integer) As UInteger
    Dim res As UInteger = 0
    res += array(offset)
    res += array(offset + 1) * CUInt(&H100)
    res += array(offset + 2) * CUInt(&H10000)
    res += array(offset + 3) * CUInt(&H1000000)
    Return res
  End Function

  Private Sub SetWavDevicePowerRequirement()
    If Me._hWavPowerReq = IntPtr.Zero Then
      Me._hWavPowerReq = SetPowerRequirement("WAV1:", DevicePowerState.D0, (POWER_NAME Or POWER_FORCE), IntPtr.Zero, 0)
    End If
  End Sub

  Private Sub ReleaseWavDevicePowerRequirement()
    If Me._hWavPowerReq <> IntPtr.Zero Then
      ReleasePowerRequirement(Me._hWavPowerReq)
      Me._hWavPowerReq = IntPtr.Zero
    End If
  End Sub

  Private Delegate Sub WriteErrorOnMainThreadDelegate(ByVal err As String)
  Public Sub WriteErrorOnMainThread(ByVal err As String)
    If MainFormReference IsNot Nothing AndAlso Settings.PowerManagementDebugLogging Then
      If MainFormReference.InvokeRequired Then
        Dim del As New WriteErrorOnMainThreadDelegate(AddressOf WriteErrorOnMainThread)
        MainFormReference.Invoke(del, New Object() {err})
      Else
        ErrorLog.WriteToErrorLog(err)
      End If
    End If
  End Sub

#Region " IDisposable Support "
  Private disposedValue As Boolean = False    ' To detect redundant calls

  ' IDisposable
  Protected Overridable Sub Dispose(ByVal disposing As Boolean)
    If Not Me.disposedValue Then
      If disposing Then
        Me.StopPreventDeviceSleep()
      End If
    End If
    Me.disposedValue = True
  End Sub

  ' This code added by Visual Basic to correctly implement the disposable pattern.
  Public Sub Dispose() Implements IDisposable.Dispose
    ' Do not change this code.  Put cleanup code in Dispose(ByVal disposing As Boolean) above.
    Dispose(True)
    GC.SuppressFinalize(Me)
  End Sub
#End Region

End Class
